import os
import sys
from urllib import request

from flask import Flask, request, send_file, jsonify

from ai_translator_v2.translator import TranslationConfig, PDFTranslator
from ai_translator_v2.utils import LOG, ArgumentParser

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

app = Flask(__name__)
TEMP_FILE_DIR = "flask_temps/"


@app.route('/translate', methods=['POST'])
def translation():
    try:
        input_file = request.files['input_file']
        source_language = request.form['source_language']
        print(request.form['source_language'])

        target_language = request.form['target_language']

        LOG.debug(f"[input_file]\n{input_file}")
        LOG.debug(f"[input_file.filename]\n{input_file.filename}")

        if input_file and input_file.filename:
            # create temp file.
            input_file_path = TEMP_FILE_DIR + input_file.filename
            LOG.debug(f"[input_file_path]\n{input_file_path}")
            input_file.save(input_file_path)
            output_file_path = Translator.translate_pdf(
                input_file_path,
                source_language=source_language,
                target_language=target_language)
            # 移除临时文件
            os.remove(input_file_path)
            # 构造完整的文件路径
            output_file_path = os.getcwd() + "/" + output_file_path
            LOG.debug(f"[output_file_path]\n{output_file_path}")
            return send_file(output_file_path, as_attachment=True)

    except Exception as e:
        response = {
            'status': 'error',
            'message': str(e)
        }
        return jsonify(response), 400


def initialize_translator():
    # 解析命令行
    argument_parser = ArgumentParser()
    args = argument_parser.parse_args()

    # 初始化配置单例
    config = TranslationConfig()
    config.initialize(args)
    # 实例化 PDFTranslator 类，并调用 translate_pdf() 方法
    global Translator
    Translator = PDFTranslator(config.model_name)


def mkdir_temp_file_dir():
    if not os.path.exists(TEMP_FILE_DIR):
        os.mkdir(TEMP_FILE_DIR)


if __name__ == "__main__":
    mkdir_temp_file_dir()
    initialize_translator()
    app.run(host='0.0.0.0', port=3000, debug=True)
